import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense

class NeuralNetworkPredictor:
    def __init__(self):
        # 加载数据
        self.df = pd.read_csv('./NN/scenic_data.csv')
        self.model = None
        
    def create_sequences(self, data, n_steps):
        """创建用于LSTM的序列数据"""
        X, y = [], []
        for i in range(len(data) - n_steps):
            X.append(data[i:i+n_steps])
            y.append(data[i+n_steps, :18])  # 预测前18个特征
        return np.array(X), np.array(y)
    
    def build_and_train_model(self, n_steps=7, epochs=50):
        """构建并训练LSTM模型"""
        # 数据预处理
        scaler = StandardScaler()
        scaled_data = scaler.fit_transform(self.df.values)
        
        # 创建序列数据
        X, y = self.create_sequences(scaled_data, n_steps)
        
        # 划分训练集和测试集
        train_size = int(len(X) * 0.8)  # 修正为0.8
        X_train, X_test = X[:train_size], X[train_size:]
        y_train, y_test = y[:train_size], y[train_size:]
        
        print(f"训练集形状: {X_train.shape}, 测试集形状: {X_test.shape}")
        
        # 构建LSTM模型
        self.model = Sequential([
            LSTM(50, activation='relu', return_sequences=True, 
                 input_shape=(n_steps, X_train.shape[2])),
            LSTM(50, activation='relu'),
            Dense(18)  # 输出层对应18个特征
        ])
        
        # 编译模型
        self.model.compile(optimizer='adam', loss='mse')
        
        # 训练模型
        history = self.model.fit(
            X_train, y_train,
            epochs=epochs,
            validation_data=(X_test, y_test),
            verbose=1
        )
        
        # 评估模型
        test_loss = self.model.evaluate(X_test, y_test, verbose=0)
        print(f"测试集损失: {test_loss:.4f}")
        
        # 保存模型
        self.model.save('NN/my_model.keras')
        print("模型已保存至: NN/my_model.keras")
        
        return history

if __name__ == '__main__':
    # 创建神经网络预测器实例
    predictor = NeuralNetworkPredictor()
    
    # 训练模型
    predictor.build_and_train_model(n_steps=7, epochs=50)